import unittest

def show(p):
    for i in range(len(p)):
        print p[i]

#DO NOT USE IMPORT
#ENTER CODE BELOW HERE
#ANY CODE ABOVE WILL CAUSE
#HOMEWORK TO BE GRADED
#INCORRECT

def calculate(colors, measurements, motions, sensor_right, p_move):    
    nr = len(colors)
    nc = len(colors[0])
    p = [[1./(nr * nc)] * nc] * nr
    for k in range(len(measurements)):
        p = move(p, nr, nc, motions[k])
        p = sense(p, nr, nc, measurements[k])
    return p

def sense(p, nr, nc, Z):
    q=[]
    for r in range(nr):
        newr=[]
        for c in range(nc):
            newr.append(p[r][c] * sensor_right if colors[r][c] == Z else p[r][c] *  (1 - sensor_right))
        q.append(newr)
    s = 0      
    for r in range(nr): 
        s = s + sum(q[r])
    for r in range(nr): 
        for c in range(nc): 
            q[r][c] = q[r][c] / s
    return q

def move(p, nr, nc, U):
    q = []
    for r in range(nr):
        newr=[]
        for c in range(nc): 
            newr.append(p[(r - U[0]) % nr][(c - U[1]) % nc] * p_move + p[r][c] * (1 - p_move))
        q.append(newr)
    return q
    
#Your probability array must be printed 
#with the following code.


# Tests

class Homework1_Tests(unittest.TestCase):

    def setUp(self):
        self.precision = 0.000001

    def _VerifyMatrix(self, expected, p):
        nr = len(p)
        nc = len(p[0]) 
        for r in range(nr):
            for c in range(nc):
                self.assertTrue(abs(expected[r][c] - p[r][c]) < self.precision, "Different at index (" + str(r) + ", " + str(c) + "). Expected: " + str(expected[r][c]) + ". Actual: " + str(p[r][c]))

    def test1(self):
        global colors
        global measurements
        global motions
        global sensor_right
        global p_move

        # SETUP
        colors = [['green', 'green', 'green'],
                  ['green',   'red', 'green'],
                  ['green', 'green', 'green']]

        measurements = ['red']
        motions = [[0,0]]
        sensor_right = 1.0
        p_move = 1.0

        # EXECUTE
        p = calculate(colors, measurements, motions, sensor_right, p_move)

        # VERIFY
        expected = [[0, 0, 0],
                    [0, 1, 0],
                    [0, 0, 0]]

        self._VerifyMatrix(expected, p)

    def test2(self):
        global colors
        global measurements
        global motions
        global sensor_right
        global p_move

        # SETUP
        colors = [['green', 'green', 'green'],
                  ['green',   'red',   'red'],
                  ['green', 'green', 'green']]

        measurements = ['red']
        motions = [[0,0]]
        sensor_right = 1.0
        p_move = 1.0

        # EXECUTE
        p = calculate(colors, measurements, motions, sensor_right, p_move)

        # VERIFY
        expected = [[0,   0,   0],
                    [0, 0.5, 0.5],
                    [0,   0,   0]]

        self._VerifyMatrix(expected, p)

    def test3(self):
        global colors
        global measurements
        global motions
        global sensor_right
        global p_move

        # SETUP
        colors = [['green', 'green', 'green'],
                  ['green',   'red',   'red'],
                  ['green', 'green', 'green']]

        measurements = ['red', 'red']
        motions = [[0,0], [0,1]]
        sensor_right = 0.8
        p_move = 1.0

        # EXECUTE
        p = calculate(colors, measurements, motions, sensor_right, p_move)

        # VERIFY
        expected = [[0.033333333, 0.033333333, 0.0333333333],
                    [0.133333333, 0.133333333, 0.5333333333],
                    [0.033333333, 0.033333333, 0.0333333333]]

        self._VerifyMatrix(expected, p)

    def test4(self):
        global colors
        global measurements
        global motions
        global sensor_right
        global p_move

        # SETUP
        colors = [['green', 'green', 'green'],
                  ['green',   'red',   'red'],
                  ['green', 'green', 'green']]

        measurements = ['red', 'red']
        motions = [[0,0], [0,1]]
        sensor_right = 1.0
        p_move = 1.0

        # EXECUTE
        p = calculate(colors, measurements, motions, sensor_right, p_move)

        # VERIFY
        expected = [[0, 0, 0],
                    [0, 0, 1],
                    [0, 0, 0]]

        self._VerifyMatrix(expected, p)

    def test5(self):
        global colors
        global measurements
        global motions
        global sensor_right
        global p_move

        # SETUP
        colors = [['green', 'green', 'green'],
                  ['green',   'red',   'red'],
                  ['green', 'green', 'green']]

        measurements = ['red', 'red']
        motions = [[0,0], [0,1]]
        sensor_right = 0.8
        p_move = 0.5

        # EXECUTE
        p = calculate(colors, measurements, motions, sensor_right, p_move)

        # VERIFY
        expected = [[0.0289855, 0.0289855, 0.0289855],
                    [0.0724637, 0.2898550, 0.4637681],
                    [0.0289855, 0.0289855, 0.0289855]]

        self._VerifyMatrix(expected, p)

    def test6(self):
        global colors
        global measurements
        global motions
        global sensor_right
        global p_move

        # SETUP
        colors = [['green', 'green', 'green'],
                  ['green',   'red',   'red'],
                  ['green', 'green', 'green']]

        measurements = ['red', 'red']
        motions = [[0,0], [0,1]]
        sensor_right = 1.0
        p_move = 0.5

        # EXECUTE
        p = calculate(colors, measurements, motions, sensor_right, p_move)

        # VERIFY
        expected = [[0, 0, 0],
                    [0, 0.333333, 0.666666],
                    [0, 0, 0]]

        self._VerifyMatrix(expected, p)

    def test7(self):
        global colors
        global measurements
        global motions
        global sensor_right
        global p_move

        # SETUP
        colors = [['red', 'green', 'green', 'red' , 'red'],
                  ['red', 'red', 'green', 'red', 'red'],
                  ['red', 'red', 'green', 'green', 'red'],
                  ['red', 'red', 'red', 'red', 'red']]

        measurements = ['green', 'green', 'green' ,'green', 'green']
        motions = [[0,0],[0,1],[1,0],[1,0],[0,1]]
        sensor_right = 0.7
        p_move = 0.8

        # EXECUTE
        p = calculate(colors, measurements, motions, sensor_right, p_move)

        # VERIFY
        expected = [[0.011059, 0.024640, 0.067996, 0.044724, 0.024651],
                    [0.007153, 0.010171, 0.086965, 0.079884, 0.009350],
                    [0.007397, 0.008943, 0.112729, 0.353507, 0.040655],
                    [0.009106, 0.007153, 0.014349, 0.043133, 0.036425]]

        self._VerifyMatrix(expected, p)

    def testAllPossibleMoves(self):
        global colors
        global measurements
        global motions
        global sensor_right
        global p_move

        # SETUP
        colors = [['red', 'green', 'green', 'red' , 'red'],
                  ['red', 'red', 'green', 'red', 'red'],
                  ['red', 'red', 'green', 'green', 'red'],
                  ['red', 'red', 'red', 'red', 'red']]

        measurements = ['red', 'red', 'red' ,'red', 'green', 'green', 'green']
        motions = [[0,0],[0,1],[1,0],[1,0],[0,-1],[0,-1],[-1,0]]
        sensor_right = 0.9
        p_move = 1.0

        # EXECUTE
        p = calculate(colors, measurements, motions, sensor_right, p_move)

        # VERIFY
        expected = [[1.3367016013355119e-05, 0.08770099206362302, 0.08770099206362303, 0.0010827282970817648, 0.00012030314412019606], 
                    [1.6502488905376682e-07, 0.00012030314412019609, 0.7893089285726073, 0.009744554673735889, 0.00012030314412019606],
                    [1.3367016013355119e-05, 1.3367016013355119e-05, 0.009744554673735887, 0.009744554673735887, 0.0010827282970817648],
                    [0.00012030314412019609, 0.0010827282970817648, 0.0010827282970817648, 0.0010827282970817648, 0.00012030314412019609]]

        self._VerifyMatrix(expected, p)

if __name__ == '__main__':
    unittest.main()
